# main.py - 多方式认证系统（Pydantic V2 实现）
from fastapi import FastAPI, Depends, HTTPException, status
from fastapi.security import OAuth2PasswordBearer, OAuth2PasswordRequestForm  # Add this line
from pymongo import MongoClient
from pymongo.collection import Collection
from bson import ObjectId
from pydantic import (
    BaseModel,
    EmailStr,
    Field,
    field_validator,
    ValidationError,
    model_validator
)
from passlib.context import CryptContext
import random
import string
from datetime import datetime, timedelta
import jwt
from typing import Union, Any, Optional

# === 系统配置 ===
app = FastAPI(
    title="Multi-Auth System",
    description="支持多方式注册/登录的认证系统（Pydantic V2 版）",
    version="1.1.0"
)

# 安全相关配置（生产环境请用环境变量）
SECRET_KEY = "your-secret-key-keep-safe"  # 必须修改！
ALGORITHM = "HS256"
ACCESS_TOKEN_EXPIRE_MINUTES = 30
SMS_CODE_EXPIRE_MINUTES = 5  # 验证码有效期

# 数据库连接（本地开发配置）
MONGO_URI = "mongodb://localhost:27017/"
DB_NAME = "auth_system"
client = MongoClient(MONGO_URI)
db = client[DB_NAME]
users: Collection = db["users"]

# 初始化索引（启动时执行）
try:
    users.drop_index("identifier_1")
except:
    pass
users.create_index(
    [("identifier", 1)],
    unique=True,
    name="identifier_unique"
)

# 安全工具
pwd_context = CryptContext(schemes=["bcrypt"], deprecated="auto")
oauth2_scheme = OAuth2PasswordBearer(tokenUrl="login")

# === 数据模型 ===
class UserCreate(BaseModel):
    """注册请求模型（自动处理标识符格式）"""
    identifier: str = Field(..., min_length=5, max_length=20)
    password: str = Field(..., min_length=8)
    identifier_type: str = "username"  # 默认为用户名类型

    @field_validator("identifier")
    @classmethod
    def normalize_identifier(cls, v: str):
        """统一处理标识符格式"""
        return v.strip()

class UserDB(BaseModel):
    """数据库存储的用户模型"""
    id: Optional[str] = None
    identifier: str
    hashed_password: str
    verification_code: Optional[str] = None
    code_expires: Optional[datetime] = None
    is_active: bool = True
    created_at: datetime = Field(default_factory=datetime.utcnow)

    @classmethod
    def from_mongo(cls, doc: dict) -> "UserDB":
        """MongoDB 文档转模型"""
        if doc and "_id" in doc:
            doc["id"] = str(doc["_id"])
            del doc["_id"]
        return cls(**doc)

    def model_dump(self, **kwargs):
        """转换为字典，排除None值"""
        data = super().model_dump(**kwargs)
        return {k: v for k, v in data.items() if v is not None}

# === 核心功能 ===
def create_access_token(subject: str, expires_delta: Optional[timedelta] = None):
    """创建 JWT 访问令牌"""
    expire = datetime.utcnow() + (expires_delta or timedelta(minutes=15))
    return jwt.encode(
        {"sub": subject, "exp": expire},
        SECRET_KEY,
        algorithm=ALGORITHM
    )

async def get_current_user(token: str = Depends(oauth2_scheme)):
    """依赖项：获取当前用户"""
    try:
        payload = jwt.decode(token, SECRET_KEY, algorithms=[ALGORITHM])
        user_doc = users.find_one({"identifier": payload["sub"]})
        if not user_doc:
            raise HTTPException(status_code=401, detail="Invalid token")
        return UserDB.from_mongo(user_doc)
    except (jwt.PyJWTError, ValidationError):
        raise HTTPException(status_code=401, detail="Invalid credentials")

# === 路由定义 ===
@app.get("/users", tags=["管理员"])
async def get_all_users(skip: int = 0, limit: int = 10):
    """获取所有用户信息（分页查询）"""
    cursor = users.find().skip(skip).limit(limit)
    return [UserDB.from_mongo(doc).model_dump() for doc in cursor]

@app.delete("/users", tags=["管理员"])
async def delete_all_users():
    """删除所有用户（仅开发环境使用）"""
    result = users.delete_many({})
    return {"message": f"Deleted {result.deleted_count} users"}

@app.get("/users/{user_id}", tags=["管理员"])
async def get_user(user_id: str, current_user: UserDB = Depends(get_current_user)):
    """获取单个用户信息（需管理员权限）"""
    # 验证用户ID格式
    if not user_id or len(user_id) != 24:
        raise HTTPException(status_code=400, detail="Invalid user ID format")
    
    # 查询用户信息
    user_doc = users.find_one({"_id": ObjectId(user_id)})
    if not user_doc:
        raise HTTPException(status_code=404, detail="User not found")
    
    return UserDB.from_mongo(user_doc).model_dump()

@app.delete("/users/{user_id}", tags=["管理员"])
async def delete_user(user_id: str, current_user: UserDB = Depends(get_current_user)):
    """删除单个用户（需管理员权限）"""
    # 验证用户ID格式
    if not user_id or len(user_id) != 24:
        raise HTTPException(status_code=400, detail="Invalid user ID format")
    
    # 检查用户是否存在
    user_doc = users.find_one({"_id": ObjectId(user_id)})
    if not user_doc:
        raise HTTPException(status_code=404, detail="User not found")
    
    # 执行删除操作
    result = users.delete_one({"_id": ObjectId(user_id)})
    if result.deleted_count == 1:
        return {"message": "User deleted successfully", "user_id": user_id}
    raise HTTPException(status_code=500, detail="Failed to delete user")

@app.post("/register", tags=["认证"], status_code=status.HTTP_201_CREATED)
async def register_user(user: UserCreate):
    """多方式注册（用户名/邮箱/手机号）"""
    try:
        # 检查唯一性
        if users.find_one({"identifier": user.identifier}):
            raise HTTPException(
                status_code=400,
                detail=f"User already registered"
            )
        
        # 存储哈希密码
        hashed_password = pwd_context.hash(user.password)
        user_db = UserDB(
            identifier=user.identifier,
            hashed_password=hashed_password
        )
        result = users.insert_one(user_db.model_dump(exclude={"id"}))
        return {
            "message": "Registration successful",
            "user": {
                "identifier": user.identifier,
                "id": str(result.inserted_id)
            }
        }
    except Exception as e:
        raise HTTPException(status_code=400, detail=str(e))

@app.post("/request-sms", tags=["认证"])
async def request_sms_code(identifier: str):
    """请求短信验证码（仅支持手机号）"""
    # 验证标识符类型
    try:
        user_create = UserCreate(identifier=identifier, password="temp")
    except ValidationError:
        raise HTTPException(
            status_code=400,
            detail="Only phone numbers are allowed for SMS verification"
        )
    if user_create.identifier_type != "phone":
        raise HTTPException(
            status_code=400,
            detail="Only phone numbers are allowed for SMS verification"
        )
    
    # 获取用户
    user_doc = users.find_one({"identifier": identifier})
    if not user_doc:
        raise HTTPException(status_code=404, detail="User not found")
    
    # 生成带时效的验证码
    code = "".join(random.choices(string.digits, k=6))
    expires = datetime.utcnow() + timedelta(minutes=SMS_CODE_EXPIRE_MINUTES)
    users.update_one(
        {"_id": user_doc["_id"]},
        {"$set": {"verification_code": code, "code_expires": expires}}
    )
    
    # 模拟短信发送（生产环境调用真实API）
    # await send_sms(identifier, f"您的验证码：{code}（{SMS_CODE_EXPIRE_MINUTES}分钟内有效）")
    return {
        "message": "SMS verification code sent",
        "expires_in": SMS_CODE_EXPIRE_MINUTES * 60,
        "hint": "Code: " + code  # 仅开发环境显示
    }

@app.post("/login", tags=["认证"])
async def login(form_data: OAuth2PasswordRequestForm = Depends()):
    """用户登录"""
    user_doc = users.find_one({"identifier": form_data.username})
    if not user_doc:
        raise HTTPException(status_code=400, detail="Invalid credentials")
    
    user = UserDB.from_mongo(user_doc)
    if not pwd_context.verify(form_data.password, user.hashed_password):
        raise HTTPException(status_code=400, detail="Invalid credentials")
    
    access_token = create_access_token(
        subject=user.identifier,
        expires_delta=timedelta(minutes=ACCESS_TOKEN_EXPIRE_MINUTES)
    )
    return {"access_token": access_token, "token_type": "bearer"}

@app.get("/me", tags=["用户"], response_model=dict)
async def get_current_user_info(current_user: UserDB = Depends(get_current_user)):
    """获取当前用户信息（需认证）"""
    return {
        "id": current_user.id,
        "identifier": current_user.identifier,
        "type": current_user.identifier_type,
        "created_at": current_user.created_at.isoformat()
    }
    
from fastapi.staticfiles import StaticFiles
from fastapi.responses import HTMLResponse
from routes import router as api_router    
app.mount("/static", StaticFiles(directory="static"), name="static")

app.include_router(api_router)



# === 启动检查（可选）===
if __name__ == "__main__":
    import uvicorn
    uvicorn.run("main:app", host="0.0.0.0", port=8011, reload=True)
